package main

import (
	"github.com/lucas-clemente/quic-go"
	"flag"
	"github.com/lucas-clemente/quic-go/internal/utils"
	"crypto/tls"
	"time"
	"crypto/rsa"
	"crypto/x509"
	"math/big"
	"encoding/pem"
	"crypto/rand"
	"log"
	mathrand "math/rand"
	"encoding/json"
	"fmt"
	"github.com/lucas-clemente/quic-go/qerr"
	"os"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/fec"
	"io"
	"strings"
	"runtime/pprof"
)


var addr = "0.0.0.0:4242"
var addrCli = "0.0.0.0:4242"
var unreliable = false
var file = ""
var deadline = 0
var rebufferingDeadline = 2*time.Millisecond
var bitrate uint16 = 500 //Kbps
var amount uint64 = 10*1000
const BATCH_READ_WRITE = 1000
const NUMBER_OF_BYTES_TO_WRITE = 1000
var desynchronization time.Duration
var fecScheme, _ = fec.NewReedSolomonFECScheme(10)
var convoFecScheme = fec.NewRandomLinearFECScheme()
var maxPathID uint8
//var fecScheme = &fec.XORFECScheme{}

var fecSchemeArg string

type resultClient struct {
	ElapsedNS int64	`json:"elapsedNS"`
	BytesReceived uint64 `json:"bytesReceived"`
	ReBufferingTimes []uint64 `json:"reBufferingTimes"`
	Bitrate uint16 `json:"bitrate"`
}

type resultServer struct {
	ElapsedNS int64	`json:"elapsedNS"`
	BytesSent uint64 `json:"bytesSent"`
	Bitrate uint16 `json:"bitrate"`
}

// hardcoded : Unreliable = true (should be false by default)

func main() {
	address := flag.String("addr", "0.0.0.0:4242", "address to bind if server or to connect to if client")
	addresscli := flag.String("addrcli", "0.0.0.0:4242", "address of the client")
	pipe := flag.String("file", "/var/www/video_pipe", "name of the pipe to open")
	unr := flag.Bool("unreliable", true, "if present, use unreliable streams")
	dl := flag.Int("deadline", 0, "Amount of time (in milliseconds) during which unreliable Stream frames should be retransmitted")
	verbose := flag.Bool("v", false, "verbose mode")
	serv := flag.Bool("s", false, "acts like a server")
	br := flag.Uint("b", 2000, "bitrate at which the server must send the data (in Kbps)")
	am := flag.Uint64("amount", 20*1000, "amount of data to send in kiloBytes")
	bfdl := flag.Int("rebufferingdeadline", 1, "represents the rebuffering deadline (the amount of time in ms after which a rebuffering is considered)")
	desync := flag.Int("desynchronization", 10, "represents the amount of time (in milliseconds) of desynchronization between the client and the server, to represent a buffer")
	cpuprofile := flag.String("cpuprofile", "", "write cpu profile `file`")
	fecSchemeFlag := flag.String("fecScheme", "rlc", "rs or rlc")
	flag.Parse()


	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}

	}

	if *verbose {
		utils.SetLogLevel(utils.LogLevelDebug)
	} else {
		utils.SetLogLevel(utils.LogLevelInfo)
	}
	fecSchemeArg = *fecSchemeFlag
	addr = *address
	addrCli = *addresscli
	unreliable = *unr
	file = *pipe
	deadline = *dl
	bitrate = uint16(*br)
	amount = *am
	rebufferingDeadline = time.Duration(*bfdl)*time.Millisecond
	desynchronization = time.Duration(*desync)*time.Millisecond
	var err error
	if *serv {
		err = server()
	} else {
		err = client()
	}
	if err != nil {
		if quicerror, ok := err.(*qerr.QuicError); ok && quicerror.ErrorCode == qerr.PeerGoingAway {
			log.Printf("%+v", quicerror)
		} else {
			os.Exit(42)
			fmt.Println("an error occured")
		}
	}
	pprof.StopCPUProfile()
	log.Printf("stop cpuprofile")

}

func client() error {
	timesRead := 1
	//session, err := quic.DialAddr(addr, &tls.Config{InsecureSkipVerify: true}, nil)
	session, err := quic.DialAddr(addr, &tls.Config{InsecureSkipVerify: true}, &quic.Config{Versions: []quic.VersionNumber{protocol.Version39,}})
	if err != nil {
		return err
	}
	if fecSchemeArg == "rs" {
		session.SetFECScheme(fecScheme)
		defer log.Printf("FECScheme = %T", fecScheme)
	} else {
		session.SetFECScheme(convoFecScheme)
		defer log.Printf("FECScheme = %T", convoFecScheme)
	}


	streamBitRate, err := session.AcceptStream()
	if err != nil {
		return err
	}
	log.Printf("1")
	log.Printf("2")

	var read [2]byte

	streamBitRate.Read(read[:])

	bitrate := arrayToUint16(read)

	durationToWait := getDurationToWaitFromBitrate(bitrate)

	stream, err := session.AcceptStream()
	if err != nil {
		return err
	}
	stream.SetReplayBufferSize(uint64(getNumberOfBytesForDesynchronization(desynchronization, uint(bitrate))))

	var amountRead uint64
	var bufferingTimes []uint64
	var byteRead [BATCH_READ_WRITE]byte
	firstRead := false
	now := time.Now()
	var elapsed int64
	defer func () {
		log.Printf("elapsed = %+v", elapsed)
		out, err := json.Marshal(resultClient{
			ElapsedNS: elapsed,
			BytesReceived: amountRead,
			ReBufferingTimes:bufferingTimes,
			Bitrate:bitrate,
		})
		if err != nil {
			fmt.Printf("ERROR\n")
		}
		fmt.Printf("%+v", string(out))
	}()
	for {
		beforeRead := time.Now()
		nBytes, err := stream.Read(byteRead[:])
		if err != nil {
			if err == io.EOF {
				elapsed = time.Now().Sub(now).Nanoseconds()
				// try to send a shutdown message to the server
				streamBitRate.SetUnreliable(false)
				for i := 0 ; i < 100 ; i++ {
					streamBitRate.Write([]byte{0xFF})
					time.Sleep(10*time.Millisecond)
				}
				// wait a bit for the server to receive it
				time.Sleep(250*time.Millisecond)
				break
			}
			return err
		}
		waitedTimeForRead := time.Now().Sub(beforeRead)
		amountRead+=uint64(nBytes)
		// do the check only if we have received already enough packets to avoid the congestion control warming
		if firstRead && waitedTimeForRead > durationToWait*BATCH_READ_WRITE + rebufferingDeadline {
			// there is a rebuffering
			log.Printf("rebuffering: %d", uint64((waitedTimeForRead - durationToWait*BATCH_READ_WRITE).Nanoseconds()))
			bufferingTimes = append(bufferingTimes, uint64((waitedTimeForRead - durationToWait*BATCH_READ_WRITE).Nanoseconds()))
		}
		if byteRead[nBytes-1] == 0xFF {
			amountRead--
			break
		}
		if firstRead && amountRead >= uint64(timesRead*5000) {
			log.Println(amountRead)
			timesRead++
		}
		if !firstRead {
			firstRead = true
		}
	}
	session.Close(nil)
	return nil
}

func server() error {
	protocol.NumberOfFecPackets = 20
	protocol.NumberOfInterleavedFECGroups = 1
	protocol.NumberOfRepairSymbols = 10

	//protocol.NumberOfFecPackets = 3
	//protocol.NumberOfInterleavedFECGroups = 10
	//protocol.NumberOfRepairSymbols = 1
	var bytesSent uint64
	//listener, err := quic.ListenAddr(addr, generateTLSConfig(), nil)
	listener, err := quic.ListenAddr(addr, generateTLSConfig(), &quic.Config{Versions: []quic.VersionNumber{protocol.Version39}})
	if err != nil {
		return err
	}
	sess, err := listener.Accept()
	if err != nil {
		return err
	}
	fecScheme, _ = fec.NewReedSolomonFECScheme(4)
	// fecScheme = &fec.XORFECScheme{}
	convoFecScheme = fec.NewRandomLinearFECScheme()
	if fecSchemeArg == "rs" {
		sess.SetFECScheme(fecScheme)
	} else {
		sess.SetFECScheme(convoFecScheme)
	}
	streamBitRate, err := sess.OpenStreamSync()
	if err != nil {
		return err
	}
	stream, err := sess.OpenStreamSync()
	if err != nil {
		return err
	}
	stream.SetUnreliable(unreliable)
	//stream.SetRetransmissionDeadline(time.Duration(deadline*1000000)) // 1000000 == time.Millisecond
	stream.SetRetransmissionDeadline(time.Duration(getDurationToWaitFromBitrate(bitrate)))
	array := uint16ToArrayBigEndian(bitrate)
	streamBitRate.Write(array[:])
	time.Sleep(1*time.Second)
	var toSend = make([]byte, NUMBER_OF_BYTES_TO_WRITE)
	durationToWait := getDurationToWaitFromBitrate(bitrate)
	now := time.Now()

	defer func () {
		log.Printf("elapsed = %+v", time.Now().Sub(now))
		out, err := json.Marshal(resultServer{
			ElapsedNS: time.Now().Sub(now).Nanoseconds(),
			BytesSent: bytesSent,
			Bitrate: bitrate,
		})
		if err != nil {
			fmt.Printf("ERROR\n")
		}
		fmt.Printf("%+v", string(out))
	}()
	for ; bytesSent < amount*1000/8;{
		before := time.Now()
		if bytesSent % 5000 ==  0 {
			log.Println(bytesSent)
		}
		packetSize := getPacket(bitrate, toSend)
		now = time.Now()
		stream.Write(toSend)
		if time.Now().Sub(now) > 1*time.Millisecond {
			log.Printf("timewritten = %+v", time.Now().Sub(now))
		}
		bytesSent+=uint64(packetSize)
		time.Sleep(durationToWait*time.Duration(NUMBER_OF_BYTES_TO_WRITE) - time.Now().Sub(before))
	}
	stream.Close()
	finalByteToRead := make([]byte, 1)
	streamBitRate.Read(finalByteToRead)

	return nil
}

func arrayToUint16(array [2]byte) uint16 {
	return uint16(array[1]) | (uint16(array[0]) << 8)
}

func uint16ToArrayBigEndian(val uint16) [2]byte{
	return [2]byte{byte(val>>8), byte(val&0xFF)}
}

// Setup a bare-bones TLS config for the server
func generateTLSConfig() *tls.Config {
	key, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		panic(err)
	}
	template := x509.Certificate{SerialNumber: big.NewInt(1)}
	template.DNSNames = []string{strings.Split(addr, ":")[0]}
	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
	if err != nil {
		panic(err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})

	tlsCert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		panic(err)
	}


	return &tls.Config{Certificates: []tls.Certificate{tlsCert}}
}

func getDurationToWaitFromBitrate(bitrate uint16) time.Duration {
	bitrateInKiloBytesPerSecond := float64(bitrate)/8.0	//250
	bitrateInBytesPerSecond := uint64(bitrateInKiloBytesPerSecond * 1000)  //250000
	timeToWaitBetweekEachByteInNanoSeconds := 1000000000.0/bitrateInBytesPerSecond
	return time.Duration(timeToWaitBetweekEachByteInNanoSeconds)
}

func getPacket(bitrate uint16, dst []byte) int {
	var i int
	for i = 0 ; i < NUMBER_OF_BYTES_TO_WRITE ; i++ {
		dst[i] = byte(mathrand.Intn(0xFF))
	}
	return NUMBER_OF_BYTES_TO_WRITE
}

func getNumberOfBytesForDesynchronization(desynchronization time.Duration, bitrateInKiloBitsPerSecond uint) uint {
	bitrateInBytesPerMilliSecond := bitrateInKiloBitsPerSecond/8
	desyncInMilliSeconds := desynchronization.Nanoseconds()/1000000
	return bitrateInBytesPerMilliSecond*uint(desyncInMilliSeconds)
}